import torch
import torch.nn as nn
import pandas as pd
from sklearn.preprocessing import MinMaxScaler


# 定义与训练相同的模型结构
class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(5, 64)
        self.dropout1 = nn.Dropout(p=0.1)
        self.fc2 = nn.Linear(64, 64)
        self.dropout2 = nn.Dropout(p=0.1)
        self.bn1 = nn.BatchNorm1d(64)
        self.fc3 = nn.Linear(64, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.dropout1(x)
        x = torch.relu(self.fc2(x))
        x = self.dropout2(x)
        x = self.bn1(x)
        x = self.fc3(x)
        return x


# 加载保存的模型参数
model = MLP()
model.load_state_dict(torch.load("./output/mlp_model_with_dropout.pth"))
model.eval()  # 设置模型为评估模式

# 加载新数据
new_data_df = pd.read_csv("./data/predict.csv")
new_inputs = new_data_df[["集水面积", "Sr", "Ks", "蒸发量", "降雨量(24h)"]].values

# 数据归一化
scaler = MinMaxScaler()
new_inputs_scaled = scaler.fit_transform(new_inputs)

# 转换为torch张量
X_new = torch.tensor(new_inputs_scaled, dtype=torch.float32)

# 使用模型进行预测
with torch.no_grad():
    new_predictions = model(X_new)

# 保存预测结果
new_data_df["预测洪峰流量"] = new_predictions.numpy()
new_data_df.to_csv("./output/new_predictions.csv", index=False)
print("Predictions saved to ./output/new_predictions.csv")
